Create app.py

app.py

@@ -0,0 +1,221 @@
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from transformers import GroundingDinoProcessor
+from modeling_grounding_dino import GroundingDinoForObjectDetection
+
+from PIL import Image, ImageDraw, ImageFont
+from itertools import cycle
+import os
+from datetime import datetime
+import gradio as gr
+
+# Load model and processor
+model_id = "fushh7/llmdet_swin_large_hf"
+model_id = "fushh7/llmdet_swin_tiny_hf"
+DEVICE = "cpu"
+
+print(f"[INFO] Using device: {DEVICE}")
+print(f"[INFO] Loading model from {model_id}...")
+
+processor = GroundingDinoProcessor.from_pretrained(model_id)
+model = GroundingDinoForObjectDetection.from_pretrained(model_id).to(DEVICE)
+model.eval()
+
+print("[INFO] Model loaded successfully.")
+
+# Pre-defined palette (extend or tweak as you like)
+BOX_COLORS = [
+    "deepskyblue", "red", "lime", "dodgerblue",
+    "cyan", "magenta", "yellow",
+    "orange", "chartreuse"
+]
+
+
+def save_cropped_images(original_image, boxes, labels, scores, output_dir="static/output_crops"):
+    """
+    Salva ogni regione ritagliata definita dalle bounding box in file separati.
+
+    :param original_image: Immagine PIL originale
+    :param boxes: Lista di bounding box [x_min, y_min, x_max, y_max]
+    :param labels: Lista di etichette per ogni box
+    :param scores: Lista di punteggi di confidenza
+    :param output_dir: Directory base dove salvare le immagini
+    :return: Lista dei percorsi dei file salvati
+    """
+    # Crea una directory con timestamp per evitare sovrascritture
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_path = os.path.join(output_dir, f"detections_{timestamp}")
+    os.makedirs(output_path, exist_ok=True)
+
+    saved_paths = []
+
+    for i, (box, label, score) in enumerate(zip(boxes, labels, scores)):
+        # Pulisci il label per usarlo nel nome del file
+        clean_label = "".join(c if c.isalnum() else "_" for c in label)
+
+        # Ritaglia la regione dall'immagine originale
+        cropped_img = original_image.crop(box)
+
+        # Crea il nome del file
+        filename = f"crop_{i}_{clean_label}_{score:.2f}.jpg"
+        filepath = os.path.join(output_path, filename)
+
+        # Salva l'immagine ritagliata
+        cropped_img.save(filepath)
+        saved_paths.append(filepath)
+
+    return saved_paths
+
+
+def draw_boxes(image, boxes, labels, scores, colors=BOX_COLORS, font_path="arial.ttf", font_size=16):
+    """
+    Draw bounding boxes and labels on a PIL Image.
+
+    :param image: PIL Image object
+    :param boxes: Iterable of [x_min, y_min, x_max, y_max]
+    :param labels: Iterable of label strings
+    :param scores: Iterable of scalar confidences (0-1)
+    :param colors: List/tuple of colour names or RGB tuples
+    :param font_path: Path to a TTF font for labels
+    :param font_size: Int size of font to use, default 16
+    :return: PIL Image with drawn boxes
+    """
+    # Ensure we can iterate colours indefinitely
+    colour_cycle = cycle(colors)
+    draw = ImageDraw.Draw(image)
+
+    # Pick a font (fallback to default if missing)
+    try:
+        font = ImageFont.truetype(font_path, size=font_size)
+    except IOError:
+        font = ImageFont.load_default(size=font_size)
+
+    # Assign a consistent colour per label (optional)
+    label_to_colour = {}
+
+    for box, label, score in zip(boxes, labels, scores):
+        # Reuse colour if label seen before, else take next from cycle
+        colour = label_to_colour.setdefault(label, next(colour_cycle))
+
+        x_min, y_min, x_max, y_max = map(int, box)
+
+        # Draw rectangle
+        draw.rectangle([x_min, y_min, x_max, y_max], outline=colour, width=2)
+
+        # Compose text
+        text = f"{label} ({score:.3f})"
+        text_size = draw.textbbox((0, 0), text, font=font)[2:]
+
+        # Draw text background for legibility
+        bg_coords = [x_min, y_min - text_size[1] - 4,
+                     x_min + text_size[0] + 4, y_min]
+        draw.rectangle(bg_coords, fill=colour)
+
+        # Draw text
+        draw.text((x_min + 2, y_min - text_size[1] - 2),
+                  text, fill="black", font=font)
+
+    return image
+
+def resize_image_max_dimension(image, max_size=4096):
+    """
+    Resize an image so that the longest side is at most max_size pixels,
+    while maintaining the aspect ratio.
+
+    :param image: PIL Image object
+    :param max_size: Maximum dimension in pixels (default: 1024)
+    :return: PIL Image object (resized)
+    """
+    width, height = image.size
+
+    # Check if resizing is needed
+    if max(width, height) <= max_size:
+        return image
+
+    # Calculate new dimensions maintaining aspect ratio
+    ratio = max_size / max(width, height)
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+
+    # Resize the image using high-quality resampling
+    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+def detect_and_draw(
+    img: Image.Image,
+    text_query: str,
+    box_threshold: float = 0.14,
+    text_threshold: float = 0.13,
+    save_crops: bool = True
+) -> Image.Image:
+    """
+    Detect objects described in `text_query`, draw boxes, return the image.
+    Note: `text_query` must be lowercase and each concept ends with a dot
+          (e.g. 'a cat. a remote control.')
+    """
+
+    # Make sure text is lowered
+    text_query = text_query.lower()
+
+    # If the image size is too large, we make it smaller
+    img = resize_image_max_dimension(img, max_size=4096)
+
+    # Preprocess the image
+    inputs = processor(images=img, text=text_query, return_tensors="pt").to(DEVICE)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    results = processor.post_process_grounded_object_detection(
+        outputs,
+        inputs.input_ids,
+        box_threshold=box_threshold,
+        text_threshold=text_threshold,
+        target_sizes=[img.size[::-1]]
+    )[0]
+
+    img_out = img.copy()
+    img_out = draw_boxes(
+        img_out,
+        boxes  = results["boxes"].cpu().numpy(),
+        labels = results.get("text_labels", results.get("labels", [])),
+        scores = results["scores"]
+    )
+    if save_crops:
+        saved_paths = save_cropped_images(
+            img,
+            boxes=results["boxes"].cpu().numpy(),
+            labels=results.get("text_labels", results.get("labels", [])),
+            scores=results["scores"]
+        )
+        print(f"Saved {len(saved_paths)} cropped images to: {os.path.dirname(saved_paths[0])}")
+
+    return img_out
+
+# Create example list
+examples = [
+    ["examples/stickers.jpg", "stickers. labels.", 0.24, 0.23],
+#    ["examples/IMG_8920.jpeg", "bin. water bottle. hand. shoe.", 0.4, 0.3],
+#    ["examples/IMG_9435.jpeg", "lettuce. orange slices (group). eggs (group). cheese (group). red cabbage. pear slices (group).", 0.4, 0.3],
+]
+
+# Create Gradio demo
+app = gr.Interface(
+    fn      = detect_and_draw,
+    inputs  = [
+        gr.Image(type="pil", label="Image"),
+        gr.Textbox(value="stickers",
+                   label="Text Query (lowercase, end each with '.', for example 'a bird. a tree.')"),
+        gr.Slider(0.0, 1.0, 0.14, 0.05, label="Box Threshold"),
+        gr.Slider(0.0, 1.0, 0.13, 0.05, label="Text Threshold")
+    ],
+    outputs = gr.Image(type="pil", label="Detections"),
+    title   = "Sticker Geo Tagger",
+    description = f"""Upload an image containings stickers and adjust thresholds to see detections.
+    <a href='/output_crops/' target='crops'>output_crops</a>
+    """,
+    examples = examples,
+    cache_examples = True,
+)
+
+#app.launch(server_name="0.0.0.0", server_port=22590, root_path="/stikkiers2", share=False)
+app.launch(server_name="0.0.0.0", share=False)