Viewing File: /home/ubuntu/combine_ai/bg_upscale/src/main.py

from removebg.remover import *
import os
import uuid

def generate_unique_filename():
    unique_id = str(uuid.uuid4())
    return unique_id

obj = BackgroundRemover()


async def img_remover(input_path, user_id, image_type = 'png'):
    try:

        output_image_name = generate_unique_filename() + f'.{image_type}'

        image_output_file_dir = os.path.join(os.environ['RMG_OUTPUT_DIRECTORY'], str(user_id))
        if not os.path.exists(image_output_file_dir):
            os.makedirs(image_output_file_dir)
        image_output_file_path = os.path.join(os.environ['RMG_OUTPUT_DIRECTORY'], str(user_id), output_image_name)

        # Create the directory if it does not exist
        live_path = os.path.join(os.environ['OUTPUT_MEDIA_URL'], 'removebg/data/user_data/output_data', str(user_id), output_image_name)
        print("live_path", live_path)

        result = obj.process(input_path, image_output_file_path)

        return {"image_output_file_path": image_output_file_path, "live_path": live_path, "message": "Success"}
    
    except Exception as e:
        print(f"Failed to remove background: {e}")
        return {"image_output_file_path": None, "live_path": None, "error": str(e), "message": "Failed"}
    

async def upscale_image(input_path, user_id, image_type = 'png'):
    try:

        output_image_name = generate_unique_filename() + f'.{image_type}'
        image_output_file_path = os.path.join(os.environ['SCALE_OUTPUT_DIRECTORY'], str(user_id), output_image_name)

        # Create the directory if it does not exist
        image_output_file_dir = os.path.join(os.environ['SCALE_OUTPUT_DIRECTORY'], str(user_id))
        if not os.path.exists(image_output_file_dir):
            os.makedirs(image_output_file_dir)
        
        live_path = os.path.join(os.environ['OUTPUT_MEDIA_URL'], 'imscale/data/user_data/output_data', str(user_id), output_image_name)
        print("live_path", live_path)

        result = await scale(input_path, image_output_file_path)

        if result:
            return {"image_output_file_path": image_output_file_path, "live_path": live_path, "message": "Success"}
        else:
            return {"image_output_file_path": None, "live_path": None, "message": "Failed"}
    
    except Exception as e:
        print(f"Failed to remove background: {e}")
        return {"image_output_file_path": None, "live_path": None, "error": str(e), "message": "Failed"}

import cv2
async def scale(input_path, output_path):
    try:
        image = cv2.imread(input_path)
        bicubic = cv2.resize(image, (image.shape[1]*2, image.shape[0]*2), interpolation=cv2.INTER_CUBIC)
        cv2.imwrite(output_path, bicubic)

        return True
    except Exception as e:
        print(f"Failed to upscale image: {e}")
        return False