Create app.py

app.py

@@ -0,0 +1,182 @@
+import gradio as gr
+import io
+import numpy as np
+import torch
+from decord import cpu, VideoReader, bridge
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import BitsAndBytesConfig
+import json
+
+# Model Configuration
+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
+
+# Define delay reasons for each step
+DELAY_REASONS = {
+    "Step 1": ["No raw material available", "Person repatching the tire"],
+    "Step 2": ["Person repatching the tire", "Lack of raw material"],
+    "Step 3": ["Person repatching the tire", "Lack of raw material"],
+    "Step 4": ["Person repatching the tire", "Lack of raw material"],
+    "Step 5": ["Person repatching the tire", "Lack of raw material"],
+    "Step 6": ["Person repatching the tire", "Lack of raw material"],
+    "Step 7": ["Person repatching the tire", "Lack of raw material"],
+    "Step 8": ["No person available to collect tire", "Person repatching the tire"]
+}
+
+def load_video(video_data, strategy='chat'):
+    bridge.set_bridge('torch')
+    mp4_stream = video_data
+    num_frames = 24
+    decord_vr = VideoReader(io.BytesIO(mp4_stream), 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():
+    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):
+    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, possible_reasons):
+    return f"""Analyze the video of Step {step_number} in the tire manufacturing process.
+    
+Possible delay reasons for this step are:
+{', '.join(possible_reasons)}
+
+Based on the video evidence, determine which of these reasons best explains the delay.
+Please provide:
+1. Your chosen reason from the list above
+2. Specific visual evidence supporting this choice
+3. Brief explanation of why other reasons are less likely
+
+Focus your analysis on visual cues that support your conclusion."""
+
+def inference(video, step_number, selected_reason):
+    if not video:
+        return "Please upload a video first."
+    
+    try:
+        model, tokenizer = load_model()
+        video_data = video.read()
+        
+        # Get possible reasons for the selected step
+        possible_reasons = DELAY_REASONS[step_number]
+        
+        # Generate the analysis prompt
+        prompt = get_analysis_prompt(step_number, possible_reasons)
+        
+        # Get model prediction
+        temperature = 0.8
+        response = predict(prompt, video_data, temperature, model, tokenizer)
+        
+        return response
+        
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+def update_reasons(step):
+    """Update the dropdown choices based on the selected step"""
+    return gr.Dropdown(choices=DELAY_REASONS[step])
+
+# Gradio Interface
+def create_interface():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                video = gr.Video(label="Upload Manufacturing Video", sources=["upload"])
+                step_number = gr.Dropdown(
+                    choices=list(DELAY_REASONS.keys()),
+                    label="Manufacturing Step",
+                    value="Step 1"
+                )
+                reason = gr.Dropdown(
+                    choices=DELAY_REASONS["Step 1"],
+                    label="Select Delay Reason",
+                    value=DELAY_REASONS["Step 1"][0]
+                )
+                analyze_btn = gr.Button("Analyze Delay", variant="primary")
+            
+            with gr.Column():
+                output = gr.Textbox(label="Analysis Result", lines=10)
+        
+        # Update reasons when step changes
+        step_number.change(
+            fn=update_reasons,
+            inputs=[step_number],
+            outputs=[reason]
+        )
+        
+        # Trigger analysis when button is clicked
+        analyze_btn.click(
+            fn=inference,
+            inputs=[video, step_number, reason],
+            outputs=[output]
+        )
+    
+    return demo
+
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()