app.py

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


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


DELAY_REASONS = {
    "Step 1": ["Delay in Bead Insertion","Lack of raw material"],
    "Step 2": ["Inner Liner Adjustment by Technician","Person rebuilding defective Tire Sections"],
    "Step 3": ["Manual Adjustment in Ply1 apply","Technician repairing defective Tire Sections"],
    "Step 4": ["Delay in Bead set","Lack of raw material"],
    "Step 5": ["Delay in Turnup","Lack of raw material"],
    "Step 6": ["Person Repairing sidewall","Person rebuilding defective Tire Sections"],
    "Step 7": ["Delay in sidewall stitching","Lack of raw material"],
    "Step 8": ["No person available to load Carcass","No person available to collect tire"]
}

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, possible_reasons):
    """Constructs the prompt for analyzing delay reasons based on the selected step."""
    return f"""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:
You are analyzing video footage from Step {step_number} of a tire manufacturing process where a delay has been detected. Your task is to determine the most likely cause of the delay from the following possible reasons:
{', '.join(possible_reasons)}
### Task Context:  
You are analyzing video footage from a specific step in the tire manufacturing process where a delay has been detected. Your task is to determine the most likely cause of the delay from the following possible reasons, while tracking the timeline of events to suggest actions that need to be taken at specific points.  

### Required Analysis:  
Analyze the delay in the movement of the following objects from the provided manufacturing video:  

1. **h_stock_left**: Identified by contours with the color **Green**.  
2. **h_stock_right**: Identified by contours with the color **Pink**.  
3. **compressor_metal**: Identified by contours with the color **Orange**.  
4. **conveyor2**: Identified by contours with the color **Blue**.  
5. **white_down_roller_left**: Identified by contours with the color **White**.  
6. **conveyor1**: Identified by contours with the color **Brown**.  

### Steps for Analysis:  
1. **Contour Detection**: Extract the contours for each specified color.  
2. **Movement Tracking**: Track the movement of each object across video frames and log timestamps where delays or anomalies occur.  
3. **Delay Identification**: Identify and measure delays or inconsistencies in their expected movement patterns.  
4. **Action Timeline**:  
   - Provide timestamps where specific actions are required based on observed delays.  
   - Suggest actions to resolve delays based on identified reasons and visual evidence.  

### Additional Observations:  
- If no person is visible in any frames, the delay reason might be due to their absence.  
- If a person is visible and is observed interacting with the tire layers, it could indicate an issue requiring patching or adjustments.  

### Analysis Framework:  
Analyze the frames and contours for objects such as:  
- **h_stock_left**, **h_stock_right**, **conveyor1**, **conveyor2**, **compressor_metal**, **person**, **orange_roller_metal_left**, **orange_roller_metal_right**, **white_down_roller_left**, **white_down_roller_right**, and **vaccum_blue**.  

Compare the observed evidence against the possible delay reasons. Select the most likely reason based on visual cues.  

### Provide Your Analysis in the Following Format:  
1. **Selected Reason**: [State the most likely reason from the given options].  
2. **Visual Evidence**: [Describe specific visual cues that support your selection].  
3. **Reasoning**: [Explain why this reason best matches the observed evidence].  
4. **Action Timeline**:  
   - [Timestamp]: [Describe the action needed].  
   - [Timestamp]: [Describe the next action needed].  
5. **Alternative Analysis**: [Brief explanation of why other possible reasons are less likely].  

### Important Notes:  
- Focus on precise and observable visual evidence from the video.  
- Provide time-based actionable insights to address detected delays.  
- Clearly state if no person or specific activity is observed.  
"""




# Load model globally
model, tokenizer = load_model()

def inference(video, step_number):
    """Analyzes video to predict the most likely cause of delay in the selected manufacturing step."""
    try:
        if not video:
            return "Please upload a video first."
        
        possible_reasons = DELAY_REASONS[step_number]
        prompt = get_analysis_prompt(step_number, possible_reasons)
        temperature = 0.8
        response = predict(prompt, video, temperature, model, tokenizer)
        
        return response
    except Exception as e:
        return f"An error occurred during analysis: {str(e)}"

def create_interface():
    """Creates the Gradio interface for the Manufacturing Delay Analysis System with examples."""
    with gr.Blocks() as demo:
        gr.Markdown("""
        # Manufacturing Delay Analysis System
        Upload a video of the manufacturing step and select the step number. 
        The system will analyze the video and determine the most likely cause of delay.
        """)
        
        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"
                )
                analyze_btn = gr.Button("Analyze Delay", variant="primary")
            
            with gr.Column():
                output = gr.Textbox(label="Analysis Result", lines=10)
        
        # Add examples
        examples = [
            ["7838_step2_2_eval.mp4", "Step 2"],
            ["7838_step6_2_eval.mp4", "Step 6"],
            ["7838_step8_1_eval.mp4", "Step 8"],
            ["7993_step6_3_eval.mp4", "Step 6"],
            ["7993_step8_3_eval.mp4", "Step 8"]
            
        ]
        
        gr.Examples(
            examples=examples,
            inputs=[video, step_number],
            cache_examples=False
        )
        
        analyze_btn.click(
            fn=inference,
            inputs=[video, step_number],
            outputs=[output]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.queue().launch(share=True)