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: \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 , (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)