SITHVID20 Contract Tracing Challenge

contact-tracing.py

from math import sqrt
from PIL import Image

def main():
    img = Image.open("contact-tracing.png")
    infected = get_color_coords(img, Ic)
    healthy = get_color_coords(img, Hc)
    contact = []

    for i in infected:
        for h in healthy:
            if distance(i,h) <= min_safe_distance:
                if h not in contact:
                    contact.append(h)

    print("Attendees that need a SITHVID-20 test (%d):" % len(contact))
    gen_contact_map(img, contact)
    print("See generated 'sithvid20-contact-map.png' image")

def gen_contact_map(img, contact):
    map = img.copy()
    for c in contact:
        map.putpixel(c, (255,255,0))
    map.save("sithvid20-contact_map.png")

def distance(coord1, coord2):
    return(sqrt((coord1[0] - coord2[0]) ** 2 + (coord1[1] - coord2[1]) ** 2))
    

def get_color_coords(img, target_color):
    x, y = img.size
    targets = []
    for i in range(0,x):
        for j in range(0,y):
            if img.getpixel((i,j)) == target_color:
                targets.append((i,j))
    return(targets)

if __name__ == "__main__":
    Hc = (0,180,230)
    Ic = (255,0,0)
    min_safe_distance = 6
    main()