aa

main.py

@@ -1,4 +1,5 @@
 from io import BytesIO
+import io
 import random
 
 import requests
@@ -54,8 +55,9 @@ def get_mask_replicate(input_pil, positive_prompt, expand_by=0, negative_prompt=
     }
 
     s3filepath = f"target/{os.urandom(20).hex()}.png"
-    input_buffer = BytesIO()
-    input_pil.save(input_buffer, 'JPEG')  # Use the appropriate format
+    input_buffer = io.BytesIO()
+    input_pil = input_pil.convert("RGB")  # Convert to RGB if it's not
+    input_pil.save(input_buffer, 'PNG')
     input_buffer.seek(0)
     s3.put_object(Bucket=S3_BUCKET_NAME, Key=s3filepath, Body=input_buffer)
 
@@ -276,6 +278,8 @@ def get_nude(original_pil, original_max_size=2000, generate_max_size=768, positi
     person_cropped_pil = crop_to_coords(
         crop_coord_1, crop_coord_2, small_original_image)
 
+    person_cropped_pil.save("verrueckt.jpg")
+
     expanded_mask_image = get_mask_replicate(
         person_cropped_pil, "bra . blouse . skirt . dress", negative_prompt="face", expand_by=10, replicate_api_key=replicate_api_key)
 

test_it.py

@@ -1,52 +1,78 @@
-from main import get_nude
-import os
+from io import BytesIO
 from PIL import Image
+import requests
+import time
+import os
+import boto3
+
+
+S3_REGION = "fra1"
+S3_ACCESS_ID = "0RN7BZXS59HYSBD3VB79"
+S3_ACCESS_SECRET = "hfSPgBlWl5jsGHa2xuByVkSpancgVeA2CVQf2EMp"
+S3_ENDPOINT_URL = "https://s3.solarcom.ch"
+S3_BUCKET_NAME = "pissnelke"
+
+s3_session = boto3.session.Session()
+s3 = s3_session.client(
+    service_name="s3",
+    region_name=S3_REGION,
+    aws_access_key_id=S3_ACCESS_ID,
+    aws_secret_access_key=S3_ACCESS_SECRET,
+    endpoint_url=S3_ENDPOINT_URL,
+)
+
+def get_mask_replicate(input_pil, positive_prompt, expand_by=0, negative_prompt="", replicate_api_key=""):
+    # Set up the API endpoint and headers
+    api_endpoint = "https://api.replicate.com/v1/predictions"
+    headers = {
+        "Authorization": f"Token {replicate_api_key}"
+    }
+
+    s3filepath = f"target/{os.urandom(20).hex()}.png"
+    input_buffer = BytesIO()
+    input_pil.save(input_buffer, 'JPEG')  # Use the appropriate format
+    input_buffer.seek(0)
+    s3.put_object(Bucket=S3_BUCKET_NAME, Key=s3filepath, Body=input_buffer)
+
+    # Prepare the data for the POST request
+    data = {
+        "version": "ee871c19efb1941f55f66a3d7d960428c8a5afcb77449547fe8e5a3ab9ebc21c",
+        "input": {
+            "image": f"{S3_ENDPOINT_URL}/{S3_BUCKET_NAME}/{s3filepath}",
+            "mask_prompt": positive_prompt,
+            "negative_mask_prompt": negative_prompt,
+            "adjustment_factor": expand_by,
+        }
+    }
+
+    # Make the initial POST request
+    response = requests.post(api_endpoint, json=data, headers=headers)
+    response_data = response.json()
+
+    print(response_data)
+
+    # Check the status of the prediction and wait for completion
+    while True:
+        prediction_response = requests.get(f"{api_endpoint}/{response_data['id']}", headers=headers)
+        prediction_data = prediction_response.json()
+
+        if prediction_data['status'] == 'failed':
+            raise Exception(prediction_data.get('error'))
+
+        if prediction_data.get('status') == 'succeeded':
+            output_link = prediction_data['output'][2]
+            break
+
+        time.sleep(1)  # Avoid spamming the server, wait for a bit before the next status check
+
+    # Get the output image
+    output_response = requests.get(output_link)
+    image_data = BytesIO(output_response.content)
+
+    # Use PIL to handle the image
+    output_image = Image.open(image_data)
+
+    return output_image
 
-def create_grid_image(original, edited1, edited2):
-    # Create a new image with appropriate size to hold the three images (original, edited1, edited2)
-    width, height = original.size
-    grid = Image.new('RGB', (width, height * 3))
-    
-    # Paste the original and edited images into the grid
-    grid.paste(original, (0, 0))
-    grid.paste(edited1, (0, height))
-    grid.paste(edited2, (0, height * 2))
-
-    return grid
-
-def process_images(input_dir, output_dir, edit_function):
-    # Create the output directory if it doesn't exist
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-
-    # Loop through all files in the input directory
-    for filename in os.listdir(input_dir):
-        print(f"Processing {filename}")
-        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
-            base_filename, file_extension = os.path.splitext(filename)
-            grid_filename = f"{base_filename}1{file_extension}"
-            
-            if os.path.exists(os.path.join(output_dir, grid_filename)):
-                print(f"Skipping {filename}, already exists in output directory")
-                continue
-            
-            # Construct full file path
-            file_path = os.path.join(input_dir, filename)
-
-            # Process the image
-            original = Image.open(file_path)
-            try:
-                edited_images = edit_function(original)
-            except Exception as e:
-                print(f"Error processing {filename}: {e}")
-                continue
-
-            # Create the grid image
-            grid_image = create_grid_image(original, edited_images[0], edited_images[1])
-
-            # Save the grid image
-            grid_image.save(os.path.join(output_dir, grid_filename))
-
-            print(f"Saved {grid_filename}")
-
-process_images('./dataset', './dataset_result_grids', get_nude)
+verrueckt_pil = Image.open("verrueckt.jpg")
+get_mask_replicate(verrueckt_pil, "person", replicate_api_key="r8_GTeyENFqfOXFAI0COiGlB2RkhqEzqS64XBuIk")