inference.py

import os
import cv2
import torch
import numpy as np
import gradio as gr
from PIL import Image, ImageDraw
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
from transformers import OwlViTProcessor, OwlViTForObjectDetection
import gc

models = {
	'vit_b': './checkpoints/sam_vit_b_01ec64.pth',
	'vit_l': './checkpoints/sam_vit_l_0b3195.pth',
	'vit_h': './checkpoints/sam_vit_h_4b8939.pth'
}

image_examples = [
    [os.path.join(os.path.dirname(__file__), "./images/53960-scaled.jpg"), 0, []],
    [os.path.join(os.path.dirname(__file__), "./images/2388455-scaled.jpg"), 1, []],
    [os.path.join(os.path.dirname(__file__), "./images/1.jpg"),2,[]],
    [os.path.join(os.path.dirname(__file__), "./images/2.jpg"),3,[]],
    [os.path.join(os.path.dirname(__file__), "./images/3.jpg"),4,[]],
    [os.path.join(os.path.dirname(__file__), "./images/4.jpg"),5,[]],
    [os.path.join(os.path.dirname(__file__), "./images/5.jpg"),6,[]],
    [os.path.join(os.path.dirname(__file__), "./images/6.jpg"),7,[]],
    [os.path.join(os.path.dirname(__file__), "./images/7.jpg"),8,[]],
    [os.path.join(os.path.dirname(__file__), "./images/8.jpg"),9,[]]
]


def plot_boxes(img, boxes):
	img_pil = Image.fromarray(np.uint8(img * 255)).convert('RGB')
	draw = ImageDraw.Draw(img_pil)
	for box in boxes:
		color = tuple(np.random.randint(0, 255, size=3).tolist())
		x0, y0, x1, y1 = box
		x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1)
		draw.rectangle([x0, y0, x1, y1], outline=color, width=6)
	return img_pil


def segment_one(img, mask_generator, seed=None):
	if seed is not None:
		np.random.seed(seed)
	masks = mask_generator.generate(img)
	sorted_anns = sorted(masks, key=(lambda x: x['area']), reverse=True)
	mask_all = np.ones((img.shape[0], img.shape[1], 3))
	for ann in sorted_anns:
		m = ann['segmentation']
		color_mask = np.random.random((1, 3)).tolist()[0]
		for i in range(3):
			mask_all[m == True, i] = color_mask[i]
	result = img / 255 * 0.3 + mask_all * 0.7
	return result, mask_all


def generator_inference(device, model_type, points_per_side, pred_iou_thresh, stability_score_thresh,
                        min_mask_region_area, stability_score_offset, box_nms_thresh, crop_n_layers, crop_nms_thresh,
                        input_x, progress=gr.Progress()):
	# sam model
	sam = sam_model_registry[model_type](checkpoint=models[model_type]).to(device)
	mask_generator = SamAutomaticMaskGenerator(
		sam,
		points_per_side=points_per_side,
		pred_iou_thresh=pred_iou_thresh,
		stability_score_thresh=stability_score_thresh,
		stability_score_offset=stability_score_offset,
		box_nms_thresh=box_nms_thresh,
		crop_n_layers=crop_n_layers,
		crop_nms_thresh=crop_nms_thresh,
		crop_overlap_ratio=512 / 1500,
		crop_n_points_downscale_factor=1,
		point_grids=None,
		min_mask_region_area=min_mask_region_area,
		output_mode='binary_mask'
	)

	# input is image, type: numpy
	if type(input_x) == np.ndarray:
		result, mask_all = segment_one(input_x, mask_generator)
		return result, mask_all
	elif isinstance(input_x, str):  # input is video, type: path (str)
		cap = cv2.VideoCapture(input_x)  # read video
		frames_num = cap.get(cv2.CAP_PROP_FRAME_COUNT)
		W, H = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
		fps = int(cap.get(cv2.CAP_PROP_FPS))
		out = cv2.VideoWriter("output.mp4", cv2.VideoWriter_fourcc('x', '2', '6', '4'), fps, (W, H), isColor=True)
		for _ in progress.tqdm(range(int(frames_num)),
		                       desc='Processing video ({} frames, size {}x{})'.format(int(frames_num), W, H)):
			ret, frame = cap.read()  # read a frame
			result, mask_all = segment_one(frame, mask_generator, seed=2023)
			result = (result * 255).astype(np.uint8)
			out.write(result)
		out.release()
		cap.release()
		return 'output.mp4'


def predictor_inference(device, model_type, input_x, input_text, selected_points, owl_vit_threshold=0.1):
	# sam model
	sam = sam_model_registry[model_type](checkpoint=models[model_type]).to(device)
	predictor = SamPredictor(sam)
	predictor.set_image(input_x)  # Process the image to produce an image embedding

	if input_text != '':
		# split input text
		input_text = [input_text.split(',')]
		print(input_text)
		# OWL-ViT model
		processor = OwlViTProcessor.from_pretrained('./checkpoints/models--google--owlvit-base-patch32')
		owlvit_model = OwlViTForObjectDetection.from_pretrained("./checkpoints/models--google--owlvit-base-patch32").to(device)
		# get outputs
		input_text = processor(text=input_text, images=input_x, return_tensors="pt").to(device)
		outputs = owlvit_model(**input_text)
		target_size = torch.Tensor([input_x.shape[:2]]).to(device)
		results = processor.post_process_object_detection(outputs=outputs, target_sizes=target_size,
		                                                  threshold=owl_vit_threshold)

		# get the box with best score
		scores = torch.sigmoid(outputs.logits)
		# best_scores, best_idxs = torch.topk(scores, k=1, dim=1)
		# best_idxs = best_idxs.squeeze(1).tolist()

		i = 0  # Retrieve predictions for the first image for the corresponding text queries
		boxes_tensor = results[i]["boxes"]  # [best_idxs]
		boxes = boxes_tensor.cpu().detach().numpy()
		# boxes = boxes[np.newaxis, :, :]
		transformed_boxes = predictor.transform.apply_boxes_torch(torch.Tensor(boxes).to(device),
		                                                          input_x.shape[:2])  # apply transform to original boxes
		# transformed_boxes = transformed_boxes.unsqueeze(0)
		print(transformed_boxes.size(), boxes.shape)
	else:
		transformed_boxes = None

	# points
	if len(selected_points) != 0:
		points = torch.Tensor([p for p, _ in selected_points]).to(device).unsqueeze(1)
		labels = torch.Tensor([int(l) for _, l in selected_points]).to(device).unsqueeze(1)
		transformed_points = predictor.transform.apply_coords_torch(points, input_x.shape[:2])
		print(points.size(), transformed_points.size(), labels.size(), input_x.shape, points)
	else:
		transformed_points, labels = None, None

	# predict segmentation according to the boxes
	masks, scores, logits = predictor.predict_torch(
		point_coords=transformed_points,
		point_labels=labels,
		boxes=transformed_boxes,  # only one box
		multimask_output=False,
	)
	masks = masks.cpu().detach().numpy()
	mask_all = np.ones((input_x.shape[0], input_x.shape[1], 3))
	for ann in masks:
		color_mask = np.random.random((1, 3)).tolist()[0]
		for i in range(3):
			mask_all[ann[0] == True, i] = color_mask[i]
	img = input_x / 255 * 0.3 + mask_all * 0.7
	if input_text != '':
		img = plot_boxes(img, boxes_tensor)  # image + mask + boxes

	# free the memory
	if input_text != '':
		owlvit_model.cpu()
		del owlvit_model
	del input_text
	gc.collect()
	torch.cuda.empty_cache()

	return img, mask_all


def run_inference(device, model_type, points_per_side, pred_iou_thresh, stability_score_thresh, min_mask_region_area,
                  stability_score_offset, box_nms_thresh, crop_n_layers, crop_nms_thresh, owl_vit_threshold, input_x,
                  input_text, selected_points):
	# if input_x is int, the image is selected from examples
	if isinstance(input_x, int):
		input_x = cv2.imread(image_examples[input_x][0])
		input_x = cv2.cvtColor(input_x, cv2.COLOR_BGR2RGB)
	if (input_text != '' and not isinstance(input_x, str)) or len(selected_points) != 0:  # user input text or points
		print('use predictor_inference')
		print('prompt text: ', input_text)
		print('prompt points length: ', len(selected_points))
		return predictor_inference(device, model_type, input_x, input_text, selected_points, owl_vit_threshold)
	else:
		print('use generator_inference')
		return generator_inference(device, model_type, points_per_side, pred_iou_thresh, stability_score_thresh,
		                           min_mask_region_area, stability_score_offset, box_nms_thresh, crop_n_layers,
		                           crop_nms_thresh, input_x)