CS50 P-Shirt

import sys
import os
from PIL import Image, ImageOps

# Step 1: Check arguments
if len(sys.argv) != 3:
    sys.exit("Too few command-line arguments" if len(sys.argv) < 3 else "Too many command-line arguments")

input_file = sys.argv[1]
output_file = sys.argv[2]

# Step 2: Validate extensions
valid_extensions = [".jpg", ".jpeg", ".png"]
input_ext = os.path.splitext(input_file)[1].lower()
output_ext = os.path.splitext(output_file)[1].lower()

if input_ext not in valid_extensions or output_ext not in valid_extensions:
    sys.exit("Invalid output")

if input_ext != output_ext:
    sys.exit("Input and output have different extensions")

# Step 3: Open images
try:
    input_image = Image.open(input_file)
except FileNotFoundError:
    sys.exit("Input does not exist")

shirt = Image.open("shirt.png")

# Step 4: Resize and overlay
input_cropped = ImageOps.fit(input_image, shirt.size)
input_cropped.paste(shirt, shirt)

# Step 5: Save result
input_cropped.save(output_file)