Assign face colors based on face location in python blender

Support Materials and Textures
,
1

I am new to blender python interface. I want to assign a color to a polygon(or face) based on its vertices locations. I have looked around and came up with the following code by referencing Setting Vertex Color via Python. However, I found that len(mesh.data.vertex_colors[0].data) != len(mesh.data.polygons). I don’t quite understand why this is the case, and how can I fix this problem.

faces = mesh.data.polygons
vertices = mesh.data.vertices
color_layer = mesh.data.vertex_colors.active
for i, f in enumerate(mesh.data.vertex_colors[0].data):
    v_ind = faces[i].vertices
    vertex1 = vertices[v[0]].co
    vertex2 = vertices[v[1]].co
    vertex3 = vertices[v[2]].co
    f.color = get_face_color(vertex1, vertex2, vertex3)

UPDATE: I found out that len(mesh.data.vertex_colors[0].data) = 3*len(mesh.polygons), does that mean the code should be something like this instead?

faces = mesh.data.polygons
vertices = mesh.data.vertices
face_loop_ind = 0
for i, f in enumerate(faces):
        flidx0 = face_loop_ind
        flidx1 = face_loop_ind+1
        flidx2 = face_loop_ind+2
        face_loop_ind += 3
        v_ind = faces[i].vertices
        mesh.data.vertex_colors[0].data[flidx0].color = get_vert_color(vertices[v_ind[0]].co)
        mesh.data.vertex_colors[0].data[flidx1].color = get_vert_color(vertices[v_ind[1]].co)
        mesh.data.vertex_colors[0].data[flidx2].color = get_vert_color(vertices[v_ind[2]].co)
2 
import bpy, bmesh

mesh = bpy.context.object.data

# create bmesh from mesh
bm = bmesh.new()
bm.from_mesh(mesh)

# write to the active vertex color layer or create a new layer, if there's none
if len(bm.loops.layers.color) > 0:
    color_layer = bm.loops.layers.color[mesh.vertex_colors.active_index]
else:
    color_layer = bm.loops.layers.color.new("Color")

# colorize verts based on world location
for vert in bm.verts:
    for loop in vert.link_loops:
        loop[color_layer] = vert.co

# write bmesh data to mesh
bm.to_mesh(mesh)

or this:
(replace the “colorize verts based on world location” section to this)

from functools import reduce

for face in bm.faces:
    median = reduce(lambda x, y: x + y, [v.co for v in face.verts]) / len(face.verts)
    for loop in face.loops:
        loop[color_layer] = median