Fixing the pipeline

.gitignore

@@ -4,4 +4,5 @@
 **/__pycache__
 .pytest_cache
 .ruff_cache
-junit
+junit
+.idea

Makefile

@@ -17,4 +17,20 @@ lint:
 .PHONY: test
 test:
 	@echo Running tests
-	@poetry run pytest -v
+	@poetry run pytest -v
+
+.PHONY: generate
+generate:
+	@echo Generating images
+	@poetry run python image_handler/scripts/generate.py
+
+.PHONY: health_check_fix
+health_check_fix:
+	@echo Checking database health
+	@poetry run python image_handler/scripts/health_check.py --fix
+
+
+.PHONY: health_check
+health_check:
+	@echo Checking database health
+	@poetry run python image_handler/scripts/health_check.py

README.md

@@ -1,9 +1,9 @@
 ### Useful commands:
 - Creating the environment: `make local-setup`
-- Entering the environment: `poetry shell`
-- Leaving the environment: `exit`
 - Installing libraries in the environment: `make install`
-- Run the Flask application: `python image_handler/demo.py`
+- Run the health check script: `make health_check`
+- Run the health check script with fixes: `make health_check_fix`
+- Run the generate images script: `make generate`
 
 Download ollama from https://ollama.com/download
 Once installed, run ollama.exe, then launch terminal and type ''' ollama run llama3'''

image_handler/constants.py

@@ -0,0 +1,17 @@
+IMAGE_URL = "http://127.0.0.1:5000/image"
+DATE_URL = "http://127.0.0.1:5000/date"
+IMAGE_VERIFICATION_URL = "http://localhost:5000/image/verification-portal"
+
+START_DATE = 1725163200
+SECONDS_PER_DAY = 86400
+
+NUM_INFERENCE_STEPS = 1
+GUIDANCE_SCALE = 7
+
+# ANSI escape codes for colors
+GREEN = "\033[32m"
+RED = "\033[31m"
+WHITE = "\033[37m"
+RESET = "\033[0m"
+
+MAX_IMAGES = 5

image_handler/db.py

@@ -0,0 +1,96 @@
+import http.client
+import io
+import json
+from http import HTTPStatus
+from typing import Union
+
+import requests
+from image_handler.constants import DATE_URL, IMAGE_URL
+from image_handler_client.schemas.image_info import ImageInfo
+from PIL import Image
+
+
+def get_date():
+    response = requests.get(
+        f"{DATE_URL}/latest",
+        timeout=10,
+    )
+    return json.loads(response.text)
+
+
+def save_image(image: Union[Image.Image, str], info: ImageInfo):
+    if isinstance(image, Image.Image):
+        image_type = "jpeg"
+        image_bytes = io.BytesIO()
+        image.save(image_bytes, format=image_type)
+        image_bytes.seek(0)
+
+        files = {"file": (info.filename, image_bytes, f"image/{image_type}")}
+
+        response = requests.post(
+            f"{IMAGE_URL}/create",
+            {
+                "real": info.real,
+                "date": info.date,
+                "theme": info.theme,
+                "status": info.status,
+            },
+            files=files,
+            timeout=10,
+        )
+    else:
+        response = requests.post(
+            f"{IMAGE_URL}/create",
+            {
+                "url": image,
+                "real": info.real,
+                "date": info.date,
+                "theme": info.theme,
+                "status": info.status,
+            },
+            files=None,
+            timeout=5,
+        )
+
+    status = response.status_code
+    if status == http.client.OK:
+        print(f"Image [{info.filename}] successfully saved")
+    else:
+        print(f"Failed to save image [{info.filename}]")
+        print(response.text)
+
+
+def get_grouped_images_by_date() -> dict:
+    response = requests.get(f"{IMAGE_URL}/all-images", timeout=10)
+    response = response.json()
+
+    date_image_map = {}
+
+    # Iterate through each image in the response
+    for image in response:
+        date = image["date"]
+
+        # If the date is not already a key, create a new list for it
+        if date not in date_image_map:
+            date_image_map[date] = []
+
+        # Append the image to the list associated with this date
+        date_image_map[date].append(image)
+
+    return date_image_map
+
+
+def delete_rejected_images(date):
+    """Send a DELETE request to delete the rejected image by date and filename."""
+    try:
+        response = requests.delete(
+            f"{DATE_URL}/delete-rejected/{date}",
+            timeout=10,
+        )
+        if response.status_code == HTTPStatus.OK:
+            print(f"Successfully deleted rejected images for date: {date}")
+        else:
+            print(f"Failed to delete rejects for day: {date}. Status code: {response.status_code}")
+            print(f"Response: {response.text}")
+    except requests.exceptions.RequestException as e:
+        print(f"Error deleting image: {e}")

image_handler/errors.py

@@ -3,3 +3,13 @@ def __init__(self, expected, received):
         super().__init__(f"Expected {expected} images, but received {received}.")
         self.expected = expected
         self.received = received
+
+
+class TooManyRequestsError(Exception):
+    def __init__(self):
+        super().__init__("Exceeded maximum number of requests.")
+
+
+class OllamaError(Exception):
+    def __init__(self):
+        super().__init__("Ensure Ollama is running locally.")

image_handler/handle_rectification.py

@@ -0,0 +1,171 @@
+from image_handler.constants import GUIDANCE_SCALE, NUM_INFERENCE_STEPS
+from image_handler.db import delete_rejected_images, save_image
+from image_handler.image_handler_instance import get_image_handler
+from image_handler.ollama_bridge import create_prompt
+from image_handler.pexel_bridge import fetch_image, resize_image, save_pexel_images
+from image_handler.util import calculate_images_to_generate, convert_to_readable_date, count_images, get_image_from_url
+from image_handler_client.schemas.image_info import ImageInfo, ImageStatus
+
+
+def handle_missing_images(date, data):
+    count = count_images(data)
+    real_count = count["real"]
+    ai_count = count["ai"]
+    generate = calculate_images_to_generate(real_count, ai_count)
+    print(
+        f"Generating {generate['ai']} AI images and "
+        f"{generate['real']} real images for date: {convert_to_readable_date(date)}"
+    )
+    generate["real"] += generate["ai"]
+    generate["ai"] = 0
+
+    real_counter = 5
+    for i in range(generate["real"]):
+        response_data = fetch_image(data[0]["theme"], 80)
+        photos_data = response_data["photos"]
+        photo_urls = [photo["src"]["original"] for photo in photos_data]
+
+        answer = "n"
+        while answer != "y":
+            image = resize_image(get_image_from_url(photo_urls[real_counter]))
+            image.show()
+            real_counter += 1
+
+            answer = input(f'accept image for theme [{data[0]["theme"]}]? y/n: ')
+            if answer == "y":
+                info = ImageInfo(
+                    filename="filename",
+                    date=date,
+                    theme=data[0]["theme"],
+                    real=True,
+                    status=ImageStatus.VERIFIED.value,
+                )
+                save_image(image, info)
+                break
+
+    for i in range(generate["ai"]):
+        image_handler_instance = get_image_handler()
+
+        prompt_dict = create_prompt(theme=data[0]["theme"])
+        print(f"Using prompt: {prompt_dict['prompt']}")
+        image_handler_instance.enqueue_prompt_to_image(
+            info=ImageInfo(
+                filename="filename",
+                date=date,
+                theme=data[0]["theme"],
+                real=False,
+                status=ImageStatus.UNVERIFIED.value,
+            ),
+            kwargs={
+                "prompt": prompt_dict["prompt"],
+                "negative_prompt": prompt_dict["negative_prompt"],
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "guidance_scale": GUIDANCE_SCALE,
+                "width": 512,
+                "height": 512,
+            },
+        )
+
+        image_handler_instance.stop_processing()
+
+
+def handle_rejected_images(date: int, data: list):
+    real_counter = 5
+
+    if not any(image["status"] == ImageStatus.REJECTED.value for image in data):
+        return
+
+    # replace rejected images
+    delete_rejected_images(date)
+
+    for image_data in data:
+        image_info = ImageInfo(**image_data)
+
+        # handle real images
+        if image_info.status == ImageStatus.REJECTED.value:
+            if image_info.real:
+                response_data = fetch_image(image_info.theme, 80)
+                photos_data = response_data["photos"]
+                photo_urls = [photo["src"]["original"] for photo in photos_data]
+
+                answer = "n"
+                while answer != "y":
+                    image = resize_image(get_image_from_url(photo_urls[real_counter]))
+                    image.show()
+                    real_counter += 1
+
+                    answer = input("accept? y/n: ")
+                    if answer == "y":
+                        info = ImageInfo(
+                            filename=image_info.filename,
+                            date=date,
+                            theme=image_info.theme,
+                            real=True,
+                            status=ImageStatus.VERIFIED.value,
+                        )
+                        save_image(image, info)
+                        break
+
+            # handle AI images
+            else:
+                image_handler_instance = get_image_handler()
+
+                prompt_dict = create_prompt(theme=image_info.theme)
+                print(f"Using prompt: {prompt_dict['prompt']}")
+                image_handler_instance.enqueue_prompt_to_image(
+                    info=ImageInfo(
+                        filename=f"{image_info.theme}_{image_info.filename.split('_')[-1]}",
+                        date=date,
+                        theme=image_info.theme,
+                        real=False,
+                        status=ImageStatus.UNVERIFIED.value,
+                    ),
+                    kwargs={
+                        "prompt": prompt_dict["prompt"],
+                        "negative_prompt": prompt_dict["negative_prompt"],
+                        "num_inference_steps": NUM_INFERENCE_STEPS,
+                        "guidance_scale": GUIDANCE_SCALE,
+                        "width": 512,
+                        "height": 512,
+                    },
+                )
+
+                image_handler_instance.stop_processing()
+
+
+def add_images(theme: str, date: int, num_ai_images: int, num_pexel_images: int, starting_number: int = 0):
+    image_handler_instance = get_image_handler()
+
+    save_pexel_images(
+        info=ImageInfo(
+            filename=f"{theme}_{starting_number}",
+            date=date,
+            theme=theme,
+            real=True,
+            status=ImageStatus.UNVERIFIED.value,
+        ),
+        num_images=num_pexel_images,
+        skip=starting_number,
+    )
+
+    for i in range(num_ai_images):
+        prompt_dict = create_prompt(theme=theme)
+        print(f"Using prompt: {prompt_dict['prompt']}")
+
+        image_handler_instance.enqueue_prompt_to_image(
+            info=ImageInfo(
+                filename=f"{theme}_{i+num_pexel_images+starting_number}",
+                date=date,
+                theme=theme,
+                real=False,
+                status=ImageStatus.UNVERIFIED.value,
+            ),
+            kwargs={
+                "prompt": prompt_dict["prompt"],
+                "negative_prompt": prompt_dict["negative_prompt"],
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "guidance_scale": GUIDANCE_SCALE,
+                "width": 512,
+                "height": 512,
+            },
+        )