Untitled

 avatar
unknown
plain_text
4 months ago
4.9 kB
2
Indexable
def classify_image(self, image_path):
    try:
        print("\n=== Starting Image Classification ===")
        
        # Check internet connection
        has_internet = self.check_internet_connection()
        
        if has_internet and self.client:
            # Find the loading screen instance
            root = next(iter(CircularProgress._instances)).winfo_toplevel()
            loading_screen = None
            for widget in root.winfo_children():
                if isinstance(widget, LoadingScreen):
                    loading_screen = widget
                    break

            # Use OpenAI API with timeout
            print("Using OpenAI API classification")
            base64_image = self.encode_image(image_path)
            
            language = getattr(root, 'LANGUAGE', 'EN')
            prompt = self.bin_config.get_ai_prompt(language)

            try:
                # Update status to show API attempt
                if loading_screen:
                    loading_screen.update_status(TRANSLATIONS[language].get('classifying_image', 'Analyzing image...'))

                # Create event for timeout handling
                response_received = threading.Event()
                api_response = [None]
                api_error = [None]

                def api_call():
                    try:
                        response = self.client.chat.completions.create(
                            model="gpt-4-vision-preview",
                            messages=[
                                {
                                    "role": "user",
                                    "content": [
                                        {"type": "text", "text": prompt},
                                        {
                                            "type": "image_url",
                                            "image_url": {
                                                "url": f"data:image/jpeg;base64,{base64_image}"
                                            }
                                        }
                                    ]
                                }
                            ],
                            timeout=10  # 10 second timeout
                        )
                        api_response[0] = response
                        response_received.set()
                    except Exception as e:
                        api_error[0] = e
                        response_received.set()

                # Start API call in separate thread
                api_thread = threading.Thread(target=api_call)
                api_thread.start()

                # Wait for response with timeout
                if not response_received.wait(15):  # 15 second timeout
                    if loading_screen:
                        loading_screen.update_status("API taking too long, switching to local model...")
                    print("API timeout, falling back to local model")
                    result = self.classify_with_local_model(image_path)
                else:
                    if api_error[0]:
                        if loading_screen:
                            loading_screen.update_status("API error, switching to local model...")
                        print(f"API error: {api_error[0]}, falling back to local model")
                        result = self.classify_with_local_model(image_path)
                    else:
                        response = api_response[0]
                        result = response.choices[0].message.content.strip()
                        print('Completion Tokens:', response.usage.completion_tokens)
                        print('Prompt Tokens:', response.usage.prompt_tokens)
                        print('Total Tokens:', response.usage.total_tokens)

            except Exception as e:
                if loading_screen:
                    loading_screen.update_status("API error, switching to local model...")
                print(f"API error: {e}, falling back to local model")
                result = self.classify_with_local_model(image_path)
        
        else:
            # Use local model
            print("Using local model classification")
            if not self.interpreter:
                raise Exception("Local model not initialized")
            result = self.classify_with_local_model(image_path)
        
        print(f"\nClassification Result: {result}")
        
        # Find the matching bin ID
        bin_id = self.bin_config.get_bin_id_by_name(result)
        if bin_id:
            print(f"\nmove {bin_id} to home")
        
        print("=== Classification Complete ===\n")
        return result
        
    except Exception as e:
        print(f"\n!!! Error during classification: {str(e)}")
        raise
    finally:
        self.cleanup_temp_files()
Editor is loading...
Leave a Comment