clipcrop.py

# Original project: https://github.com/Vishnunkumar/clipcrop/blob/main/clipcrop/clipcrop.py
import os.path
import sys

import cv2
import numpy
import numpy as np
import torch
from PIL import Image
from clip import clip

import modules.paths
from modules import shared, modelloader


def clip_boxes(boxes, shape):
    # Clip boxes (xyxy) to image shape (height, width)
    if isinstance(boxes, torch.Tensor):  # faster individually
        boxes[:, 0].clamp_(0, shape[1])  # x1
        boxes[:, 1].clamp_(0, shape[0])  # y1
        boxes[:, 2].clamp_(0, shape[1])  # x2
        boxes[:, 3].clamp_(0, shape[0])  # y2
    else:  # np.array (faster grouped)
        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, shape[1])  # x1, x2
        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0])  # y1, y2


def find_position(parent: Image, child: Image):
    w = child.width
    h = child.height
    res = cv2.matchTemplate(np.array(parent), np.array(child), cv2.TM_CCOEFF_NORMED)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    top_left = max_loc
    center_x = top_left[0] + (w / 2)
    center_y = top_left[1] + (h / 2)
    return center_x, center_y


class CropClip:
    def __init__(self):
        # Model
        model_name = 'yolov5m6v7.pt'
        model_url = 'https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m6.pt'
        model_dir = os.path.join(modules.paths.models_path, "yolo")
        model_path = modelloader.load_models(model_dir, model_url, None, '.pt', model_name)
        was_safe_unpickle = shared.cmd_opts.disable_safe_unpickle
        shared.cmd_opts.disable_safe_unpickle = True
        self.model = torch.hub.load('ultralytics/yolov5', 'custom', model_path[0])
        self.clip = None
        self.preprocess = None
        shared.cmd_opts.disable_safe_unpickle = was_safe_unpickle
        # Prevent BLIP crossfire breakage
        try:
            del sys.modules['models']
        except:
            pass

    def get_center(self, image: Image, prompt: str):
        # Load image into YOLO parser
        results = self.model(image)  # includes NMS
        # Crop each image result to an array
        cropped = results.crop(False)
        l = []
        for crop in cropped:
            l.append(Image.fromarray(crop["im"]))
        if len(l) == 0:
            l = [image]
        device = shared.device
        # Take out cropped YOLO images, and get the features?
        if not self.model or not self.preprocess:
            self.clip, self.preprocess = clip.load("ViT-B/32", device=device)
        images = torch.stack([self.preprocess(im) for im in l]).to(device)
        with torch.no_grad():
            image_features = self.clip.encode_image(images)
            image_features /= image_features.norm(dim=-1, keepdim=True)

        image_features.cpu().numpy()
        image_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).cuda()
        image_std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).cuda()

        images = [self.preprocess(im) for im in l]
        image_input = torch.tensor(np.stack(images)).cuda()
        image_input -= image_mean[:, None, None]
        image_input /= image_std[:, None, None]
        with torch.no_grad():
            image_features = self.clip.encode_image(image_input).float()
        image_features /= image_features.norm(dim=-1, keepdim=True)

        def similarity_top(similarity_list, N):
            results = zip(range(len(similarity_list)), similarity_list)
            results = sorted(results, key=lambda x: x[1], reverse=True)
            top_images = []
            scores = []
            for index, score in results[:N]:
                scores.append(score)
                top_images.append(l[index])
            return scores, top_images

        # @title Crop
        with torch.no_grad():
            # Encode and normalize the description using CLIP
            text_encoded = self.clip.encode_text(clip.tokenize(prompt).to(device))
            text_encoded /= text_encoded.norm(dim=-1, keepdim=True)

        # Retrieve the description vector and the photo vectors
        similarity = text_encoded.cpu().numpy() @ image_features.cpu().numpy().T
        similarity = similarity[0]
        scores, imgs = similarity_top(similarity, N=3)
        max_area = 0
        out = None
        for img in imgs:
            img_area = img.width * img.height
            if img_area > max_area:
                max_area = img_area
                out = img

        if not out:
            out = image
        res = cv2.matchTemplate(numpy.array(image), numpy.array(out), cv2.TM_SQDIFF)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
        # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
        top_left = min_loc
        bottom_right = (top_left[0] + out.width, top_left[1] + out.height)
        return [top_left[0], bottom_right[0], top_left[1], bottom_right[1]]

    def unload(self):
        if self.model is not None:
            self.model.to('cpu')
        if self.clip:
            self.clip.to('cpu')
        if self.preprocess:
            self.preprocess.to('cpu')

    def load(self):
        if self.model is not None:
            self.model.to(shared.device)
        if self.clip:
            self.clip.to(shared.device)
        if self.preprocess:
            self.preprocess.to(shared.device)