app.py

import os
import cv2
import json
from ultralytics import YOLO
import matplotlib.pyplot as plt
import torch
from PIL import Image
from transformers import AutoProcessor
import requests
# from huggingface_hub import hf_hub_download
# os.environ["KERAS_BACKEND"] = "jax"

#import keras
import numpy as np
#import tensorflow as tf # Or from keras.utils import load_img, img_to_array
import threading
import subprocess
from io import BytesIO
import glob
import requests

import gradio as gr

description_model_id = "google/paligemma-3b-mix-224"
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
#description_model = PaliGemmaForConditionalGeneration.from_pretrained(description_model_id)
# Load the model in bfloat16 to save memory
description_model = PaliGemmaForConditionalGeneration.from_pretrained(
    description_model_id,
    torch_dtype=torch.bfloat16,
    revision="bfloat16",
).to(device)
description_processor = AutoProcessor.from_pretrained(description_model_id)




print(f"Loading vision model: {description_model} (this should be much faster)...")

print("--- Checking storage usage ---")
os.system("du -sh ~/.cache") # Check the size of the main cache folder
print("----------------------------")
import shutil


from ultralytics.nn.tasks import DetectionModel
from ultralytics.nn.modules.conv import Conv
from ultralytics.nn.modules.block import C3, C2f, C3k2
from torch.nn import (
    Sequential,
    Conv2d,
    BatchNorm2d,
    SiLU, # Also known as Swish activation
    ReLU,
    MaxPool2d,
    Module,
    Linear,
    Dropout,
    Upsample
)
trusted_classes = [
    # Ultralytics classes
    DetectionModel,
    Conv,
    C3,
    C2f,
    C3k2,
    # Standard PyTorch Layers
    Sequential,
    Conv2d,
    BatchNorm2d,
    SiLU,
    ReLU,
    MaxPool2d,
    Module,
    Linear,
    Dropout,
    Upsample
]

# Add all trusted classes at once
torch.serialization.add_safe_globals(trusted_classes)

from huggingface_hub import login
hf_token = os.getenv("HF_TOKEN")
if hf_token:
  login(token=hf_token, add_to_git_credential=False)
  print("Hugging Face login successful!")
else:
  print("HF_TOKEN secret not found. Please add it to your Space's settings.")

# IMPORTING MODELS


#repo_car = "thirarbi/vehicle_model"
#car_file = "vehicle_model.keras"
#car_model_path = hf_hub_download(repo_id=repo_car, filename=car_file)

#repo_ethnicity = "thirarbi/ethnicity"
#ethnicity_file = "ethnicity_model.keras"
#ethnicity_model_path = hf_hub_download(repo_id=repo_ethnicity, filename=ethnicity_file)

#repo_gender = "thirarbi/gender"
#gender_file = "gender_model.keras"
#gender_model_path = hf_hub_download(repo_id=repo_gender, filename=gender_file)



# AVOID USING KERAS FOR NOW
try:
    yolo_model = YOLO('yolo11n.pt') # Using the base model as an example
    # general_model = YOLO('/content/yolov8n_plus_fruitsnveggies.pt')
    phone_model = YOLO('phone_model.pt')
    #car_model = keras.saving.load_model(car_model_path)
    motorcycle_model = YOLO('motorcycle_model.pt')
    fruitsnveggy_model = YOLO('best_fruitveggy.pt')
    # Person.keras
    #ethnicity_model = keras.saving.load_model(ethnicity_model_path)
    #gender_model = keras.saving.load_model(gender_model_path)
    #age_model = coming soon

    yolo_model.to(device)
except Exception as e:
    print(f"Error loading model: {e}")
    exit()

# CLASS NAMES (OBJECT CLASSIFICATION MODELS)
# car_class_names = [
#     'Audi', 'Bentley', 'Benz', 'BMW', 'Cadillac',
#     'Dodge', 'Ferrari', 'Ford', 'Ford Mustang', 'Kia',
#     'Lamborghini', 'Lexus', 'Maserati', 'Porsche', 'Rolls-Royce',
#     'Tesla', 'Toyota', 'Alfa Romeo', 'Hyundai'
# ]
# ethnicity_class_names = [
#     'White', 'Black', 'Asian', 'Indian', 'Mexican'
# ]
# gender_class_names = [
#     'Male', 'Female'
# ]

# def preprocess_person_image(cropped_image):
#   """
#   Preprocesses a cropped person image for the ethnicity and gender models.

#   Args:
#     cropped_image: A PIL Image object of the cropped person.

#   Returns:
#     A numpy array of the preprocessed image with shape (1, 48, 48, 1).
#   """
#   # Convert to grayscale
#   gray_image = cropped_image.convert('L')
#   # Resize to 48x48
#   resized_image = gray_image.resize((48, 48))
#   # Convert to numpy array
#   img_array = tf.keras.utils.img_to_array(resized_image)
#   # Normalize pixel values (assuming the models were trained with normalized data)
#   img_array = img_array / 255.0
#   # Add batch dimension
#   img_array = np.expand_dims(img_array, 0)
#   # Add channel dimension for grayscale (explicitly 1)
#   img_array = np.expand_dims(img_array, -1)
#   return img_array

# def car_pred(car_img_array):
#     # 2. Make the Prediction
#     img = car_img_array.resize((IMG_WIDTH, IMG_HEIGHT))
#     img_array = tf.keras.utils.img_to_array(img)
#     img_array = img_array / 255.0
#     img_array = np.expand_dims(img_array, 0) # Create batch dimension

  # 3. Decode the Prediction
  # We find the index with the highest score and use it to get the brand name.
    # predictions = car_model.predict(img_array)
    # prediction_scores = predictions[0]
    # predicted_index = np.argmax(prediction_scores)
    # predicted_class_name = car_class_names[predicted_index]
    # confidence_score = 100 * np.max(prediction_scores)

    # return predicted_class_name
    #print(f"Confidence: {confidence_score:.2f}%")

# def person_pred(person_img_array):
#   # Preprocess the cropped person image
#   preprocessed_img = preprocess_person_image(person_img_array)

#   # 2. Make the Prediction
#   ethnicity_predictions = ethnicity_model.predict(preprocessed_img)
#   ethnicity_prediction_scores = ethnicity_predictions[0]

#   gender_predictions = gender_model.predict(preprocessed_img)
#   gender_prediction_scores = gender_predictions[0]

#   # 3. Decode the Prediction
#   # We find the index with the highest score and use it to get the brand name.
#   ethnicity_predicted_index = np.argmax(ethnicity_prediction_scores)
#   ethnicity_predicted_class_name = ethnicity_class_names[ethnicity_predicted_index]
#   ethnicity_confidence_score = 100 * np.max(ethnicity_prediction_scores)

#   gender_predicted_index = np.argmax(gender_prediction_scores)
#   gender_predicted_class_name = gender_class_names[gender_predicted_index]
#   gender_confidence_score = 100 * np.max(gender_prediction_scores)

#   return f"{ethnicity_predicted_class_name} {gender_predicted_class_name}"
  #print(f"Confidence: {confidence_score:.2f}%")


def encode_image(filepath):
    """Encodes an image to a base64 string."""
    with open(filepath, "rb") as image_file:
        return image_file.read()


def final_description(raw_image, text_prompt):
    try:
        # prompt_template = f"USER: <image>\n{text_prompt}\nASSISTANT:"
        # inputs = description_processor(prompt_template, images=raw_image, return_tensors="pt").to("cuda") # Ensure it runs on GPU
        # output = description_model.generate(**inputs, max_new_tokens=200, do_sample=False)
        # response = description_processor.decode(output[0], skip_special_tokens=True)
        # assistant_response = response.split("ASSISTANT:")[-1].strip()
        inputs = description_processor(text=text_prompt, images=raw_image, return_tensors="pt").to(device)
        input_token_length = inputs.input_ids.shape[1]
        
        with torch.no_grad():
            output = description_model.generate(**inputs, max_new_tokens=1000, do_sample=False)
        
        new_tokens = output[0, input_token_length:]
        cleaned_caption = description_processor.decode(new_tokens, skip_special_tokens=True)
        
        return cleaned_caption.strip()

    except Exception as e:
        return f"Sorry, I was unable to generate a description for this part of the image: {e}"
#     # Process
#     with Image.open(image_file) as img:
#         with BytesIO() as buffer:
#             img.save(buffer, format='PNG')
#             image_bytes = buffer.getvalue()

#     full_response = ''
#     # Generate a description of the image
#     response = ollama.chat(
#         model = 'llava',
#         messages=
#          [
#             {
#                 'role': 'user',
#                 'content': text_prompt,
#                 'images': [encode_image(img_path)],
#             },
#         ]
#     )
#     return response['message']['content']
    #=======================



# DETECTION & LOCALIZATION ==========================================
input_folder = "input_folder"
output_folder = "output_folder"
os.makedirs(output_folder, exist_ok=True)
if os.path.exists(output_folder):
    shutil.rmtree(output_folder)
os.makedirs(output_folder, exist_ok=True)


def process_image(input_image, user_input, min_words):

    if input_image is None:
        return None, "Error: Please upload an image.", None

    # The model expects images in RGB format
    image = input_image.convert("RGB")
    
    results = yolo_model(image) # GENERAL OBJECT DETECTION
    annotated_img = results[0].plot()
    specialized_captions_list = []
    specialized_objects_list = []

    detections = []
    for i, box in enumerate(results[0].boxes): # Neat percentage
        cls = yolo_model.names[int(box.cls)]
        conf = round(float(box.conf) * 100, 2)

      # BOXING
        x1, y1, x2, y2 = map(float, box.xyxy[0])
        x1, y1, x2, y2 = round(x1, 2), round(y1, 2), round(x2, 2), round(y2, 2)
        detections.append({
            "label": cls,
            "confidence (%)": f"{conf}%",
            "bbox": {
                "x1": x1,
                "y1": y1,
                "x2": x2,
                "y2": y2
            }
        })
      
        # CROPPING
        # Open the image file inside the loop for each detected object
        cropped = image.crop((int(x1), int(y1), int(x2), int(y2)))
        json_output = json.dumps(detections, indent=4) #json plotting file
        cropped_filename = f"{cls}_{i+1}.png"
        output_img_path = os.path.join(output_folder, cropped_filename)
        cropped.save(output_img_path)
        specialized_objects_list.append(f'[{i+1}] {cls}')


    # SPECIALIZATION FILTER
    # Save the cropped image with a unique filename
    # if cls == 'person': # DONE
    #     specialized_captions_list.append(f'[{i+1}] {final_description(output_img_path, "Please describe the race and gender of this person, their outfit, their emotion, what are they doing, and other details about it, summarize it with 15 - 25 words.")}') # join list of specialized objects caption
    #     specialized_objects_list.append(cls)
    # elif cls == 'car': # DONE
    #     specialized_captions_list.append(f'[{i+1}] {final_description(output_img_path, f"Describe this car, its color, its brand and model, the details about it, with less than 15 words.")}')
    #     specialized_objects_list.append(cls)

    # elif cls == 'motorcycle': # YOLO # JUST RETREIVE CLASS NAME.
    #     results = motorcycle_model(cropped)
    #     for box in results[0].boxes:
    #         class_id = int(box.cls[0])
    #         class_name = motorcycle_model.names[class_id]  # Get the class name using the class ID
    #     #specialized_captions_list.append(f'[{i+1}] {specialized_caption(f'{class_name} motorcycle', cropped)}.')
    #     specialized_captions_list.append(f'[{i+1}] {final_description(output_img_path, f"Describe the {class_name} motorcycle, its color, and the details about it, with less than 15 words.")}')
    #     #specialized_captions_list.append(f'[{i+1}] {final_description(output_img_path, f"Describe this motorcycle, its brand and model, its color, and the details about it, with less than 15 words.")}')
    #     specialized_objects_list.append(class_name)

    # elif cls == 'cell phone': # YOLO
    #   results = phone_model(cropped)
    #   for box in results[0].boxes:
    #       class_id = int(box.cls[0])
    #       class_name = phone_model.names[class_id]  # Get the class name using the class ID
    #   specialized_captions_list.append(f'[{i+1}] {final_description(output_img_path, f"Describe the {class_name} smartphone, its color, and the details about it, with less than 15 words.")}')
    #   specialized_objects_list.append(class_name)

    # elif cls in ('banana', 'apple', 'carrot', 'broccoli'): #fruitnvegetables # YOLO
    #   results = fruitsnveggy_model(cropped)
    #   for box in results[0].boxes:
    #       class_id = int(box.cls[0])
    #       class_name = fruitsnveggy_model.names[class_id]  # Get the class name using the class ID
    #   specialized_captions_list.append(f'[{i+1}] {final_description(output_img_path, f"Describe the vegetable/vegetables or fruit/fruits that are in the picture. What it is, how its presented, its details, with less than 15 words.")}')
    #  specialized_objects_list.append(class_name)

    # else:
    #   specialized_captions_list.append(f'[{i+1}] {final_description(output_img_path, f"Describe the object ({class_name}) in this picture. Describe any details about in in less than 15 words.")}')

  #combined_captions = f"\n".join(specialized_captions_list) # list of llama(aided) specialized captions
    combined_captions = f"\n".join(specialized_objects_list) # list of llama(aided) specialized captions 
    torch.cuda.empty_cache()
    general_prompt = f"""
    You're an expert image analyst. Provide a highly detailed, long-form description of this image <image>, (include text you see in the image).
    Overall Summary: A brief, one-paragraph overview of the entire scene, with addition of the following request: {user_input}
    Ensure that your total description is at least {min_words} words
    """
#     2. **Hint:** The objects that are included in the pictures are: {combined_captions}
    #2. **Foreground Elements:** Describe the objects, including the skies/clouds, people, or elements closest to the viewer in meticulous detail. Mention their appearance, textures, and any actions they are performing.
#     3. **Details:** If theres people, describe their race, outfits, expression, and what they're doing. If there are vehicles/phones, describe their models and brands. If there are vegetables/fruits, describe how they're presented, and their quantity.
#     4. **Midground and Background:** Detail the environment behind the main subject. Describe the setting, landscape, architecture, and any other distant elements.
#     5. **Atmosphere and Mood:** Analyze the image's lighting, color palette, and overall composition to describe the mood or feeling it evokes (e.g., joyful, somber, peaceful, chaotic).
    

    general_caption = final_description(input_image, general_prompt) # general caption
    caption = f'List of Detected Objects: \n{combined_captions}\n================== GENERAL CAPTION ================== \n{general_caption}'


    return Image.fromarray(annotated_img), caption, json_output

iface = gr.Interface(
    fn=process_image,
    inputs=[gr.Image(type="pil", label="Upload an image!"),
            gr.Textbox(label="How do you want the image to be described?", value="As detailed as possible."),
            gr.Textbox(label="Minimum words?", value="300")],
    outputs=[gr.Image(label="Annotated Image"),
             gr.Textbox(label="Description of the image"),
             gr.JSON(label="Detected Objects (JSON)")],
    title="Image Object Detection & Captioning",
    description="Upload an image to detect objects, view bounding boxes, and get captions."
)

iface.launch(share=True)