Update app.py

app.py

@@ -6,7 +6,7 @@ from decord import cpu, VideoReader, bridge
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import BitsAndBytesConfig
 
-MODEL_PATH = "THUDM/cogvlm2-video-llama3-chat"
+MODEL_PATH = "THUDM/cogvlm2-llama3-caption"
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 TORCH_TYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 else torch.float16
 
@@ -22,8 +22,91 @@ DELAY_REASONS = {
     "Step 8": ["No person available to load Carcass", "No person available to collect tire"]
 }
 
-def get_step_info(step_number):
-    """Returns detailed information about a manufacturing step."""
+def load_video(video_data, strategy='chat'):
+    """Loads and processes video data into a format suitable for model input."""
+    bridge.set_bridge('torch')
+    num_frames = 24
+    
+    if isinstance(video_data, str):
+        decord_vr = VideoReader(video_data, ctx=cpu(0))
+    else:  
+        decord_vr = VideoReader(io.BytesIO(video_data), ctx=cpu(0))
+    
+    frame_id_list = []
+    total_frames = len(decord_vr)
+    timestamps = [i[0] for i in decord_vr.get_frame_timestamp(np.arange(total_frames))]
+    max_second = round(max(timestamps)) + 1
+    
+    for second in range(max_second):
+        closest_num = min(timestamps, key=lambda x: abs(x - second))
+        index = timestamps.index(closest_num)
+        frame_id_list.append(index)
+        if len(frame_id_list) >= num_frames:
+            break
+
+    video_data = decord_vr.get_batch(frame_id_list)
+    video_data = video_data.permute(3, 0, 1, 2)
+    return video_data
+
+def load_model():
+    """Loads the pre-trained model and tokenizer with quantization configurations."""
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=TORCH_TYPE,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4"
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        torch_dtype=TORCH_TYPE,
+        trust_remote_code=True,
+        quantization_config=quantization_config,
+        device_map="auto"
+    ).eval()
+    
+    return model, tokenizer
+
+def predict(prompt, video_data, temperature, model, tokenizer):
+    """Generates predictions based on the video and textual prompt."""
+    video = load_video(video_data, strategy='chat')
+    
+    inputs = model.build_conversation_input_ids(
+        tokenizer=tokenizer,
+        query=prompt,
+        images=[video],
+        history=[],
+        template_version='chat'
+    )
+    
+    inputs = {
+        'input_ids': inputs['input_ids'].unsqueeze(0).to(DEVICE),
+        'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to(DEVICE),
+        'attention_mask': inputs['attention_mask'].unsqueeze(0).to(DEVICE),
+        'images': [[inputs['images'][0].to(DEVICE).to(TORCH_TYPE)]],
+    }
+    
+    gen_kwargs = {
+        "max_new_tokens": 2048,
+        "pad_token_id": 128002,
+        "top_k": 1,
+        "do_sample": False,
+        "top_p": 0.1,
+        "temperature": temperature,
+    }
+    
+    with torch.no_grad():
+        outputs = model.generate(**inputs, **gen_kwargs)
+        outputs = outputs[:, inputs['input_ids'].shape[1]:]
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return response
+
+def get_analysis_prompt(step_number):
+    """Constructs the prompt for analyzing delay reasons based on the selected step."""
+    
+    # Step details dictionary included directly in the prompt
     step_details = {
         1: {
             "Name": "Bead Insertion",
@@ -127,101 +210,14 @@ def get_step_info(step_number):
             "Video_substeps_expected": {
                 "0-3 seconds": "Technician unloads(removes) carcass(tire) from the machine."
             },
-            "Potential_Delay_reasons": [
+            "Potential_Delay_Reasons": [
                 "Person not available in time(in 3 sec) to remove carcass.",
                 "Person is doing bead(ring) insertion before carcass unload causing unload to be delayed by more than 3 sec"
             ]
         }
     }
     
-    return step_details.get(step_number, {"Error": "Invalid step number. Please provide a valid step number."})
-
-
-
-def load_video(video_data, strategy='chat'):
-    """Loads and processes video data into a format suitable for model input."""
-    bridge.set_bridge('torch')
-    num_frames = 24
-    
-    if isinstance(video_data, str):
-        decord_vr = VideoReader(video_data, ctx=cpu(0))
-    else:  
-        decord_vr = VideoReader(io.BytesIO(video_data), ctx=cpu(0))
-    
-    frame_id_list = []
-    total_frames = len(decord_vr)
-    timestamps = [i[0] for i in decord_vr.get_frame_timestamp(np.arange(total_frames))]
-    max_second = round(max(timestamps)) + 1
-    
-    for second in range(max_second):
-        closest_num = min(timestamps, key=lambda x: abs(x - second))
-        index = timestamps.index(closest_num)
-        frame_id_list.append(index)
-        if len(frame_id_list) >= num_frames:
-            break
-
-    video_data = decord_vr.get_batch(frame_id_list)
-    video_data = video_data.permute(3, 0, 1, 2)
-    return video_data
-
-def load_model():
-    """Loads the pre-trained model and tokenizer with quantization configurations."""
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=TORCH_TYPE,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4"
-    )
-    
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_PATH,
-        torch_dtype=TORCH_TYPE,
-        trust_remote_code=True,
-        quantization_config=quantization_config,
-        device_map="auto"
-    ).eval()
-    
-    return model, tokenizer
-
-def predict(prompt, video_data, temperature, model, tokenizer):
-    """Generates predictions based on the video and textual prompt."""
-    video = load_video(video_data, strategy='chat')
-    
-    inputs = model.build_conversation_input_ids(
-        tokenizer=tokenizer,
-        query=prompt,
-        images=[video],
-        history=[],
-        template_version='chat'
-    )
-    
-    inputs = {
-        'input_ids': inputs['input_ids'].unsqueeze(0).to(DEVICE),
-        'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to(DEVICE),
-        'attention_mask': inputs['attention_mask'].unsqueeze(0).to(DEVICE),
-        'images': [[inputs['images'][0].to(DEVICE).to(TORCH_TYPE)]],
-    }
-    
-    gen_kwargs = {
-        "max_new_tokens": 2048,
-        "pad_token_id": 128002,
-        "top_k": 1,
-        "do_sample": False,
-        "top_p": 0.1,
-        "temperature": temperature,
-    }
-    
-    with torch.no_grad():
-        outputs = model.generate(**inputs, **gen_kwargs)
-        outputs = outputs[:, inputs['input_ids'].shape[1]:]
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return response
-
-def get_analysis_prompt(step_number):
-    """Constructs the prompt for analyzing delay reasons based on the selected step."""
-    step_info = get_step_info(step_number)
+    step_info = step_details.get(step_number, {"Error": "Invalid step number. Please provide a valid step number."})
     
     if "Error" in step_info:
         return step_info["Error"]
@@ -230,45 +226,53 @@ def get_analysis_prompt(step_number):
     standard_time = step_info["Standard Time"]
     video_substeps = step_info["Video_substeps_expected"]
     potential_delay_reasons = step_info["Potential_Delay_Reasons"]
-    
-    # Format substeps for the prompt
-    substeps_text = "\n".join([f"  {time}: {description}" for time, description in video_substeps.items()])
-    
-    # Format potential delay reasons
-    potential_reasons_text = "\n  - ".join(potential_delay_reasons)
-    
+
     return f"""
-You are an AI expert system specialized in analyzing manufacturing processes and identifying production delays in tire manufacturing. Your task is to accurately determine the specific delay reason based on video analysis.
+You are an AI expert system specialized in analyzing manufacturing processes and identifying production delays in tire manufacturing. Your role is to accurately classify delay reasons based on visual evidence from production line footage.
+Task Context:
+The following are the details of all steps in the tire manufacturing process:
+{step_details}
 
-### Step Information:
-- **Step Number**: {step_number}
-- **Step Name**: {step_name}
-- **Standard Time**: {standard_time}
+You are analyzing video footage from Step {step_number} of this process. The step is called '{step_name}', and its standard time is {standard_time}.
+Required Analysis:
+Carefully observe the video for visual cues indicating production interruption.
 
-### Expected Video Substeps:
-{substeps_text}
+Step Details:
+- **Name:** {step_name}
+- **Standard Time:** {standard_time}
+- **Video Substeps Expected:** 
+{video_substeps}
 
-### Potential Delay Reasons:
-  - {potential_reasons_text}
+Possible Delay Reasons:
+- {', '.join(potential_delay_reasons)}
 
-### Task:
-Carefully analyze the video footage for visual cues indicating production interruptions. Follow these guidelines:
-1. If no person is visible in the frames, the delay might be due to their absence.
-2. If a person is visible modifying or adjusting materials, consider related delay reasons.
-3. Base your decision on observable evidence and match it to the listed potential delay reasons.
+Analysis Instructions:
+1. Analyze the video frame by frame to identify evidence of delay or irregular activity.
+2. If no person is visible in any of the frames, the reason might be the absence of required personnel.
+3. If a person is visible and modifying tire components, it may indicate repair or alignment issues.
+4. Match the observed evidence with the possible delay reasons listed above.
 
-### Required Output Format:
-Please respond with the most likely delay reason based on the analysis in the following format:
+Output Format:
+1. **Selected Reason:** [State the most likely reason from the list above]
+2. **Visual Evidence:** [Describe specific visual cues that support your selection]
+3. **Reasoning:** [Explain why this reason best matches the observed evidence]
+4. **Alternative Analysis:** [Briefly explain why other potential reasons are less likely]
 
-**Output**:
-- **Selected Delay Reason**: [State the most probable delay reason]
-- **Visual Evidence**: [Describe specific observations from the video that support your decision]
-- **Reasoning**: [Explain why this delay reason best matches the observed evidence]
-- **Alternative Analysis**: [Briefly explain why other reasons are less likely]
+Important:
+Base your analysis solely on visual evidence from the video. Focus on concrete, observable details rather than assumptions. Clearly state if no person or specific activity is observed.
+
+Example Output:
+Output = {{
+    "Selected Reason": "Delay in bead insertion",
+    "Visual Evidence": "Technician is not present during the bead alignment process.",
+    "Reasoning": "The absence of the technician caused a delay in starting the bead insertion.",
+    "Alternative Analysis": "No raw material issues were visible, and the machine appeared functional."
+}}
 """
 
 
 
+
 model, tokenizer = load_model()
 
 def inference(video, step_number):