Update project_model.py

project_model.py

@@ -1,5 +1,7 @@
 # project_module.py
-import torch, cv2, os, time
+
+# Import libraries for ML, CV, NLP, audio, and TTS
+import torch, cv2, os
 import numpy as np
 from PIL import Image
 from ultralytics import YOLO
@@ -10,96 +12,55 @@ from huggingface_hub import login
 # Authenticate to Hugging Face using environment token
 login(token=os.environ["HUGGING_FACE_HUB_TOKEN"])
 
-# Set device for computation
+# Set device for computation (GPU if available)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load models
-yolo_model = YOLO("yolov9c.pt")
+# Load all models
+yolo_model = YOLO("yolov9c.pt")  # YOLOv9 for object detection
 depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large").to(device).eval()
 depth_feat = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
 
+# Whisper for audio transcription
 whisper_pipe = pipeline(
     "automatic-speech-recognition",
     model="openai/whisper-small",
     device=0 if torch.cuda.is_available() else -1
 )
 
+# GEMMA for image+text to text QA
 gemma_pipe = pipeline(
-    "image-text-to-text",
+    "image-to-text",
     model="google/gemma-3-4b-it",
     device=0 if torch.cuda.is_available() else -1,
     torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
 )
 
+# Text-to-speech
 tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC")
 
 # -------------------------------
-# Smart Session Management
+# Session Management Class
 # -------------------------------
 
 class VisualQAState:
-    TIMEOUT_SECONDS = 60
-
     def __init__(self):
         self.current_image: Image.Image = None
         self.visual_context: str = ""
         self.message_history = []
-        self.last_interaction_time = time.time()
-
-    def _check_timeout(self):
-        if time.time() - self.last_interaction_time > self.TIMEOUT_SECONDS:
-            self.soft_reset()
-
-    def soft_reset(self):
-        self.message_history = [{
-            "role": "user",
-            "content": [
-                {"type": "image", "image": self.current_image},
-                {"type": "text", "text": self.visual_context}
-            ]
-        }]
-        print("🔄 Session timed out: soft reset applied.")
 
     def reset(self, image: Image.Image, visual_context: str):
         self.current_image = image
         self.visual_context = visual_context
-        self.message_history = [{
-            "role": "user",
-            "content": [
-                {"type": "image", "image": self.current_image},
-                {"type": "text", "text": self.visual_context}
-            ]
-        }]
-        self.last_interaction_time = time.time()
+        self.message_history = []
 
     def add_question(self, question: str):
-        self._check_timeout()
-        if not self.message_history or self.message_history[-1]["role"] == "assistant":
-            self.message_history.append({
-                "role": "user",
-                "content": [{"type": "text", "text": question}]
-            })
-            self.last_interaction_time = time.time()
+        self.message_history.append({"role": "user", "content": question})
 
     def add_answer(self, answer: str):
-        self._check_timeout()
-        self.message_history.append({
-            "role": "assistant",
-            "content": [{"type": "text", "text": answer}]
-        })
-        self.last_interaction_time = time.time()
-
-    def export_transcript(self) -> str:
-        transcript = []
-        for turn in self.message_history:
-            role = turn["role"].capitalize()
-            for entry in turn["content"]:
-                if entry["type"] == "text":
-                    transcript.append(f"{role}: {entry['text']}")
-        return "\n\n".join(transcript)
+        self.message_history.append({"role": "assistant", "content": answer})
 
 # -------------------------------
-# Image Context Generation
+# Generate Context from Image
 # -------------------------------
 
 def generate_visual_context(pil_image: Image.Image) -> str:
@@ -121,6 +82,7 @@ def generate_visual_context(pil_image: Image.Image) -> str:
         x1, y1, x2, y2 = map(int, box.xyxy[0])
         label = class_names[int(box.cls[0])]
         conf = float(box.conf[0])
+
         depth_crop = depth_map_resized[y1:y2, x1:x2]
         avg_depth = float(depth_crop.mean()) if depth_crop.size > 0 else None
 
@@ -144,7 +106,7 @@ def generate_visual_context(pil_image: Image.Image) -> str:
     return "In the image, " + ", ".join(descriptions) + "."
 
 # -------------------------------
-# Main Processing
+# Main Multimodal Processing Function
 # -------------------------------
 
 session = VisualQAState()
@@ -166,8 +128,10 @@ def process_inputs(
 
     session.add_question(question)
 
-    gemma_output = gemma_pipe(text=session.message_history, max_new_tokens=200)
-    answer = gemma_output[0]["generated_text"][-1]["content"]
+    prompt = f"{session.visual_context}\n\nUser Question: {question}"
+
+    gemma_output = gemma_pipe(prompt, max_new_tokens=200)
+    answer = gemma_output[0]["generated_text"]
 
     session.add_answer(answer)
 
@@ -178,4 +142,3 @@ def process_inputs(
         output_audio_path = None
 
     return answer, output_audio_path
-