Update app.py

app.py

@@ -2,25 +2,22 @@ import gradio as gr
 import torch
 import numpy as np
 from transformers import OwlViTProcessor, OwlViTForObjectDetection
-from torchvision import transforms
 from PIL import Image, ImageDraw
 import cv2
 import torch.nn.functional as F
 import tempfile
-import os
-from SuperGluePretrainedNetwork.models.matching import Matching
-from SuperGluePretrainedNetwork.models.utils import read_image
 import matplotlib.pyplot as plt
 import matplotlib.cm as cm
 from io import BytesIO
+from SuperGluePretrainedNetwork.models.matching import Matching
+from SuperGluePretrainedNetwork.models.utils import read_image
 
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load models
-mixin = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
+model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32").to(device)
 processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
-model = mixin.to(device)
 
 matching = Matching({
     'superpoint': {'nms_radius': 4, 'keypoint_threshold': 0.005, 'max_keypoints': 1024},
@@ -28,14 +25,6 @@ matching = Matching({
 }).eval().to(device)
 
 # Utility functions
-def preprocess_image(image):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    return transform(image).unsqueeze(0)
-
 def save_array_to_temp_image(arr):
     rgb_arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
     img = Image.fromarray(rgb_arr)
@@ -45,24 +34,7 @@ def save_array_to_temp_image(arr):
     img.save(temp_file_name)
     return temp_file_name
 
-def stitch_images(images):
-    if not images:
-        return Image.new('RGB', (100, 100), color='gray')
-
-    max_width = max([img.width for img in images])
-    total_height = sum(img.height for img in images)
-
-    composite = Image.new('RGB', (max_width, total_height))
-
-    y_offset = 0
-    for img in images:
-        composite.paste(img, (0, y_offset))
-        y_offset += img.height
-
-    return composite
-
 def unified_matching_plot2(image0, image1, kpts0, kpts1, mkpts0, mkpts1, color, text, path=None, show_keypoints=False, fast_viz=False, opencv_display=False, opencv_title='matches', small_text=[]):
-    # Resize images to have the same height
     height = min(image0.shape[0], image1.shape[0])
     image0_resized = cv2.resize(image0, (int(image0.shape[1] * height / image0.shape[0]), height))
     image1_resized = cv2.resize(image1, (int(image1.shape[1] * height / image1.shape[0]), height))